from typing import Optional
import numpy as np
from transformers import PreTrainedModel
from transformers.trainer_callback import TrainerCallback


class ProgressiveLayerActivationCallback(TrainerCallback):
    def __init__(
        self,
        n_layers: int,
        interval_steps: int,
        model: PreTrainedModel,
        lisa_layers_attribute: Optional[str] = None,
        progression_phases: int = 4,  # 分为几个渐进阶段
        overlap_ratio: float = 0.3,   # 相邻阶段的重叠比例
        random_ratio: float = 0.2,    # 每个阶段内随机选择的比例
    ):
        super().__init__()
        self.n_layers = n_layers
        self.interval_steps = interval_steps
        self.model = model
        self.progression_phases = progression_phases
        self.overlap_ratio = overlap_ratio
        self.random_ratio = random_ratio

        # Determine the way to access layers based on the model type
        class_to_layers_map = {
            "LlamaForCausalLM": "model.model.layers",
            "Qwen2ForCausalLM": "model.model.layers", 
            "MistralForCausalLM": "model.model.layers",
            "MixtralForCausalLM": "model.model.layers",
            "GemmaForCausalLM": "model.model.layers",
            "GPT2LMHeadModel": "model.transformer.h",
            "HymbaForCausalLM": "model.model.layers",
        }
        model_class_name = self.model.__class__.__name__
        if model_class_name in class_to_layers_map:
            self.layers_attribute = class_to_layers_map[model_class_name]
        else:
            assert lisa_layers_attribute is not None, "Please provide the attribute to access the layers of the model."
            self.layers_attribute = lisa_layers_attribute
        self.total_layers = len(
            eval("self." + self.layers_attribute)
        )

        self.active_layers_indices = []
        self.current_phase = 0
        
        # 预计算每个阶段的层范围
        self.phase_ranges = self._compute_phase_ranges()
        print(f"PLA initialized with {self.progression_phases} phases:")
        for i, (start, end) in enumerate(self.phase_ranges):
            print(f"  Phase {i}: layers {start}-{end-1}")

    def _compute_phase_ranges(self):
        """计算每个渐进阶段的层范围"""
        ranges = []
        layers_per_phase = self.total_layers // self.progression_phases
        overlap_size = int(layers_per_phase * self.overlap_ratio)
        
        for phase in range(self.progression_phases):
            start = max(0, phase * layers_per_phase - overlap_size)
            end = min(self.total_layers, (phase + 1) * layers_per_phase + overlap_size)
            ranges.append((start, end))
            
        return ranges

    def freeze_all_layers(self):
        layers = eval("self." + self.layers_attribute)
        for layer in layers:
            for param in layer.parameters():
                param.requires_grad = False

    def on_step_begin(self, args, state, control, **kwargs):
        if state.global_step % self.interval_steps == 0:
            self.switch_active_layers(state.global_step)

    def switch_active_layers(self, global_step=0):
        """
        渐进式层激活策略：
        1. 根据训练进度确定当前阶段
        2. 在当前阶段的层范围内选择活跃层
        3. 结合顺序选择和随机选择，平衡稳定性和多样性
        """
        # 先冻结所有层
        self.freeze_all_layers()
        
        # 根据训练步数确定当前阶段（可以根据实际训练步数调整）
        estimated_total_steps = 1000  # 可以从args中获取或预估
        phase_duration = estimated_total_steps // self.progression_phases
        self.current_phase = min(global_step // phase_duration, self.progression_phases - 1)
        
        # 获取当前阶段的层范围
        start_layer, end_layer = self.phase_ranges[self.current_phase]
        available_layers = list(range(start_layer, end_layer))
        
        if len(available_layers) <= self.n_layers:
            # 如果可用层数不足，直接使用所有可用层
            selected_layers = available_layers
        else:
            # 混合策略：部分顺序选择 + 部分随机选择
            n_random = int(self.n_layers * self.random_ratio)
            n_sequential = self.n_layers - n_random
            
            # 顺序选择：优先选择当前阶段的核心层
            core_start = start_layer + (end_layer - start_layer - n_sequential) // 2
            sequential_layers = list(range(core_start, core_start + n_sequential))
            
            # 随机选择：从剩余层中随机选择
            remaining_layers = [l for l in available_layers if l not in sequential_layers]
            if len(remaining_layers) >= n_random:
                random_layers = np.random.choice(remaining_layers, n_random, replace=False).tolist()
            else:
                random_layers = remaining_layers
                
            selected_layers = sequential_layers + random_layers
        
        self.active_layers_indices = sorted(selected_layers)
        
        # 激活选中的层
        layers = eval("self." + self.layers_attribute)
        for idx in self.active_layers_indices:
            for param in layers[idx].parameters():
                param.requires_grad = True
                
        print(f"Step {global_step}, Phase {self.current_phase}: "
              f"Activating layers {self.active_layers_indices} "
              f"(range: {start_layer}-{end_layer-1})", flush=True)